# https://link.springer.com/chapter/10.1007/978-3-030-97281-3_14
# https://link.springer.com/content/pdf/10.1007/978-3-030-97281-3_14.pdf

import torch
from torch import nn
import einops
from axial_attention import AxialAttention, AxialPositionalEmbedding


class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden_dim, dim), nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)


class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)


class FNetBlock(nn.Module):
    def __init__(self):
        super().__init__()

    @torch.cuda.amp.autocast(enabled=False)
    def forward(self, x):
        # x = torch.fft.fftn(x, dim=(-1)).real
        # x = torch.fft.fft(torch.fft.fft(x, dim=-1), dim=-2).real
        # x = torch.pow(torch.abs(x)+1e-8, 0.8)  # abs
        return x


class FNet(nn.Module):
    def __init__(self, dim, depth, mlp_dim, dropout=0.5):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(
                nn.ModuleList([PreNorm(dim, FNetBlock()), PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))])
            )

    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x1 = ff(x) * x
            x = x + x1
        return x


class FourierPatchEmbedding(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        # (10, 64, 24, 32, 40)
        return einops.rearrange(x, 'b c z x y -> b c (z x y)')


class FourierPatchDebedding(nn.Module):
    def __init__(self, emb_shape):
        super().__init__()
        self.emb_shape = emb_shape

    def forward(self, x):
        return einops.rearrange(x, f'b c (z x y) -> b c z x y', x=self.emb_shape[1], y=self.emb_shape[2])


if __name__ == '__main__':
    x = torch.randn(1, 32, 24, 64, 80)
    x = einops.rearrange(x, 'b c z h w -> b (c h w)')
    model = FNet(x.shape[-1], 2, 64)
    y = model(x)
    import ipdb

    ipdb.set_trace()  # ! debug yusongli
